{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"../common/rfsoc_book_banner.jpg\" alt=\"University of Strathclyde\" align=\"left\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook Set C\n",
    "\n",
    "---\n",
    "\n",
    "## 05 - Fast Fourier Transform\n",
    "The previous two notebooks introduced Fourier's Theorem and the DFT. Throughout this notebook, we will investigate a technique that speeds up the computation of the DFT, known as the Fast Fourier Transform (FFT). We will begin by understanding the Danielson-Lanczos Lemma and then implement an FFT design from first principles that will serve as an example. We will briefly explore the Numpy FFT module, before finishing with a short introduction to the Short-Time Fourier Transform (STFT).\n",
    "\n",
    "## Table of Contents\n",
    "* [1. Introduction](#introduction)\n",
    "* [2. The Fast Fourier Transform](#fft)\n",
    "    * [2.1. Danielson-Lanczos Lemma](#dll)\n",
    "    * [2.2. Decimation in Time](#dit)\n",
    "    * [2.3. An FFT Example](#fft-example)\n",
    "    * [2.4. FFT Versus DFT](#fft-dft)\n",
    "    * [2.5. The Numpy FFT Module](#numpy-fft-module)\n",
    "* [3. Short-Time Fourier Transform](#stft)\n",
    "    * [3.1. STFT of a Chirp](#stft-of-a-chirp)\n",
    "* [4. Conclusion](#conclusion)\n",
    "\n",
    "## References\n",
    "* [1] - [J. W. Cooley and J. W. Tukey, \"An Algorithm for the Machine Calculation of Complex Fourier Series,\" Mathematics of Computation, Vol. 19, No. 90, 1965, pp. 297-301.](https://doi.org/10.2307/2003354)\n",
    "* [2] - [G. C. Danielson and C. Lanczos, \"Some improvements in practical Fourier analysis and their application to X-ray scattering from liquids,\" Journal of the Franklin Institute, Vol. 233, Issue 4, 1942, pp. 365-380.](https://doi.org/10.1109/ACSSC.2006.354945)\n",
    "\n",
    "## Revision\n",
    "* **v1.0** | 20/01/23 | *First Revision*\n",
    "\n",
    "---\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Introduction <a class=\"anchor\" id=\"introduction\"></a>\n",
    "To begin our investigation into the FFT, we will first redefine the DFT function from the [previous notebook](04_investigating_the_dft.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def dft(x):\n",
    "    N = np.size(x)\n",
    "    X = np.zeros(N, dtype=np.complex64)\n",
    "    n = np.arange(0, N)\n",
    "    for k in range(N):\n",
    "        X[k] = np.sum(x*np.exp(-2j*np.pi*k*n/N))\n",
    "    return X"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the purpose of plotting, we will also reintroduce the `stem_plot` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def stem_plot(x, y, title, xlabel, ylabel, xticks=None, yticks=None, figsize=(6, 3), subplot=(1, 1), bottom=None, style=None):\n",
    "    fig = plt.figure(figsize=figsize)\n",
    "    for idx, value in enumerate(y):\n",
    "        axes = fig.add_subplot(subplot[0], subplot[1], idx+1)\n",
    "        if bottom is not None:\n",
    "            axes.stem(x[idx], y[idx], use_line_collection = True, bottom=bottom[idx])\n",
    "        else:\n",
    "            axes.stem(x[idx], y[idx], use_line_collection = True)\n",
    "        if style is not None:\n",
    "            if style[idx] == 'dashed':\n",
    "                axes.plot(x[idx], y[idx], linestyle='dashed')\n",
    "        axes.grid(True, which='major')\n",
    "        if xticks is not None:\n",
    "            axes.set_xticks(xticks[idx])\n",
    "        if yticks is not None:\n",
    "            axes.set_yticks(yticks[idx])\n",
    "        axes.set_title(title[idx])\n",
    "        axes.set_xlabel(xlabel[idx])\n",
    "        axes.set_ylabel(ylabel[idx])\n",
    "        plt.box(False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we will generate a discrete waveform to be used as test stimulus when investigating the FFT. The discrete wave will have 8 sample points, a sampling frequency of 48kHz, and consist of 6kHz and 12kHz cosine waves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fs = 48e3 # Sampling frequency\n",
    "ts = 1/fs # Regular sampling period\n",
    "N = 8    # Number of samples\n",
    "k = n = np.arange(N) # Sample points\n",
    "\n",
    "\n",
    "a = 0.5*np.cos(2*np.pi*6e3*n*ts + np.pi/4)\n",
    "b = np.cos(2*np.pi*12e3*n*ts)\n",
    "\n",
    "x = np.around(a + b, 4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, use the `stem_plot` function to plot the discrete wave."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stem_plot(x      = [n],\n",
    "          y      = [x],\n",
    "          figsize= (6, 3),\n",
    "          title  = [\"Discrete Waveform $x(n)$\"], \n",
    "          style  = [\"dashed\"], \n",
    "          xlabel = [\"Samples (n)\"], \n",
    "          ylabel = [\"Amplitude\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. The Fast Fourier Transform <a class=\"anchor\" id=\"fft\"></a>\n",
    "Large DFTs consume a sizeable amount of processing power to implement. In 1965, Cooley and Tukey [1] described an efficient implementation of the DFT that we now name the Fast Fourier Transform, or FFT. The FFT exploits redundancy in the DFT to reduce the number of arithmetic operations required to obtain the frequency spectra of a discrete time waveform.\n",
    "\n",
    "In comparison to the DFT, the FFT is approximately 10 times faster and produces the exact same answer.\n",
    "\n",
    "### 2.1. Danielson-Lanczos Lemma <a class=\"anchor\" id=\"dll\"></a>\n",
    "To begin, we will explore the Danielson-Lanczos Lemma [2]. Their work describes a technique to compute an N-point DFT by separating it into two smaller size DFTs. This DFT technique is only possible if $N$ is an even number of samples (as we separate the DFT into two DFTs).\n",
    "\n",
    "The expression below separates the input waveform between two DFTs. One DFT receives all of the even samples of $x(n)$, and the other DFT is allocated the odd samples $x(n)$. Notice that $k$ is still operational across the full range of digital frequencies i.e. $k=0,1,2,...,N-1$.\n",
    "\n",
    "$$\n",
    "X(k) = \\sum_{n=0}^{N/2-1}x(2n)e^{-j2\\pi (2n)k/N} + \\sum_{n=0}^{N/2-1}x(2n+1)e^{-j2\\pi (2n+1)k/N}\n",
    "$$\n",
    "\n",
    "We can rearrange the expression above such that the complex exponential in the odd DFT can be moved outside of the summation, as below.\n",
    "\n",
    "$$\n",
    "X(k) = \\sum_{n=0}^{N/2-1}x(2n)e^{-j2\\pi (2n)k/N} + e^{-j2\\pi k/N}\\sum_{n=0}^{N/2-1}x(2n+1)e^{-j2\\pi (2n)k/N}\n",
    "$$\n",
    "\n",
    "Now, we can simply denote the same complex exponential as\n",
    "\n",
    "$$\n",
    "W_N^k = e^{-j2\\pi k/N},\n",
    "$$\n",
    "\n",
    "which is commonly referred to as the twiddle factor. The twiddle factor has a symmetric-periodic property across $N/2$ values, where\n",
    "\n",
    "$$\n",
    "W_N^k = -W_N^{k+N/2}.\n",
    "$$\n",
    "\n",
    "We can demonstrate this property in the code cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Wk = np.round(np.exp(-2j*np.pi*k/N), 2)\n",
    "\n",
    "all(Wk[0:N//2] == -Wk[N//2:N])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can rewrite the even and odd DFTs given above as $X^{even}(k)$ and $X^{odd}(k)$, respectively. The expression to compute an N-point DFT using two smaller DFTs now becomes,\n",
    "\n",
    "$$\n",
    "X(k) = X^{\\text{even}}(k) + W_N^k X^{\\text{odd}}(k).\n",
    "$$\n",
    "\n",
    "We can express the Danielson-Lanczos Lemma graphically as shown in Figure 1. The figure shown an 8-point DFT implemented using two 4-point DFTs and four multipliers.\n",
    "\n",
    "<figure>\n",
    "<img src=\"images/fft_a.png\" style=\"width: 75%\"/> \n",
    "    <figcaption><b>Figure 1: An 8-point DFT that has been implemented using two 4-point DFTs and four multipliers.</b></figcaption>\n",
    "</figure> \n",
    "\n",
    "We can implement an 8-point DFT in software using two smaller 4-point DFTs. Let us start by separating the input waveform into even and odd samples. Also, precompute the twiddle factor for values of $k$ in the range $[0, N/2-1]$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xeven = x[0::2]\n",
    "xodd = x[1::2]\n",
    "\n",
    "W8k = np.exp(-2j*np.pi*k[0:N//2]/N)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, perform each 4-point DFT on the even and odd samples. Compute the multiplication of the twiddle factor and odd samples and appropriately add or subtract the result from the even samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Xfft = np.zeros(N, dtype=np.complex64)\n",
    "\n",
    "Xeven = dft(xeven)\n",
    "Xodd = dft(xodd)\n",
    "\n",
    "temp = W8k*Xodd\n",
    "Xfft[0:N//2] = Xeven + temp\n",
    "Xfft[N//2:N] = Xeven - temp"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the results for inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Xfftm = np.abs(Xfft)\n",
    "\n",
    "stem_plot(x      = [n[0:N//2+1]*fs/N],\n",
    "          y      = [Xfftm[0:N//2+1]],\n",
    "          xticks = [n[0:N//2+1]*fs/N],\n",
    "          figsize= (6, 3),\n",
    "          title  = [\"Magnitude Spectra of $X(k)$\"], \n",
    "          xlabel = [\"Frequency (Hz)\"], \n",
    "          ylabel = [\"Amplitude\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resulting magnitude spectra of $X(k)$ appears to be correct as it clearly shows the 6kHz and 12kHz frequency components. We can double-check our result by comparing it to the DFT below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = dft(x)\n",
    "all(np.round(Xfft, 2) == np.round(X, 2))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have demonstrated that separating the DFT into two small DFTs is possible and produces the exact same result. We can now progress further by exchanging the 4-point DFTs for two 2-point DFTs in the next section."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2. Decimation in Time <a class=\"anchor\" id=\"dit\"></a>\n",
    "The Danielson-Lanczos Lemma can be applied recursively to remove the two 4-point DFTs in the previous section. This design is presented in Figure 2. The input sequence to the design is further separated into even and odd samples, which is an FFT technique known as decimate-in-time.\n",
    "\n",
    "<figure>\n",
    "<img src=\"images/fft_b.png\" style=\"width: 100%\"/> \n",
    "    <figcaption><b>Figure 2: An 8-point DFT that has been implemented using four 2-point DFTs and eight multipliers.</b></figcaption>\n",
    "</figure> \n",
    "\n",
    "As before, we can implement the design above using Python. We begin by separating the input waveform into even and odd sequences. Then, we compute the twiddle factors for $W_4^k$ and $W_8^k$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xeven_0 = x[0::4]\n",
    "xodd_0  = x[2::4]\n",
    "xeven_1 = x[1::4]\n",
    "xodd_1  = x[3::4]\n",
    "\n",
    "W8k = np.exp(-2j*np.pi*k[0:N//2]/N)\n",
    "W4k = np.exp(-2j*np.pi*k[0:N//4]/(N//2))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we apply a 2-point DFT to each input sequence. The odd sequencies are appropriately multiplied by the twiddle factor after each stage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialise registers\n",
    "Xfft = np.zeros(N, dtype=np.complex64)\n",
    "Xeven = np.zeros(N//2, dtype=np.complex64)\n",
    "Xodd = np.zeros(N//2, dtype=np.complex64)\n",
    "\n",
    "# Apply four 2-point DFTs\n",
    "Xeven_0 = dft(xeven_0)\n",
    "Xodd_0 = dft(xodd_0)\n",
    "Xeven_1 = dft(xeven_1)\n",
    "Xodd_1 = dft(xodd_1)\n",
    "\n",
    "# Get Xeven\n",
    "temp = W4k*Xodd_0\n",
    "Xeven[0:N//4] = Xeven_0 + temp\n",
    "Xeven[N//4:N//2] = Xeven_0 - temp\n",
    "\n",
    "# Get Xodd\n",
    "temp = W4k*Xodd_1\n",
    "Xodd[0:N//4] = Xeven_1 + temp\n",
    "Xodd[N//4:N//2] = Xeven_1 - temp\n",
    "\n",
    "# Obtain final output\n",
    "temp = W8k*Xodd\n",
    "Xfft[0:N//2] = Xeven + temp\n",
    "Xfft[N//2:N] = Xeven - temp"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compare the result from the decimate-in-time design and the DFT below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all(np.round(Xfft, 2) == np.round(X, 2))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An additional step for this 8-point DFT would be to apply the Danielson-Lanczos Lemma one last time to reduce the DFT into its final form. This design is the 8-point FFT implementation of the Cooley & Tukey algorithm and is given in Figure 3. It is worth while mentioning that the FFT is only able to process discrete waveforms that are a power of 2 in length.\n",
    "\n",
    "<figure>\n",
    "<img src=\"images/fft_c.png\" style=\"width: 100%\"/> \n",
    "    <figcaption><b>Figure 3: An 8-point DFT that has been implemented using twelve multipliers.</b></figcaption>\n",
    "</figure> "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3. An FFT Example <a class=\"anchor\" id=\"fft-example\"></a>\n",
    "The FFT can be implemented using a recursive function to process any power of 2 sized discrete waveform. We have implemented a function below that you can study on your own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fft(x):\n",
    "    N = np.size(x) # Must be a power of 2\n",
    "    if N == 1:\n",
    "        return x\n",
    "    else:\n",
    "        Xeven = fft(x[0::2]) # Recursion\n",
    "        Xodd = fft(x[1::2])  # Recursion\n",
    "        k = np.arange(N)\n",
    "        X = np.zeros(N, dtype=np.complex64)\n",
    "        Wk = np.exp(-2j*np.pi*k[0:N//2]/(N))\n",
    "        temp = Wk * Xodd\n",
    "        X[0:N//2] = Xeven + temp\n",
    "        X[N//2:N] = Xeven - temp\n",
    "        return X"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the FFT function above to compute the spectrum of $x(n$ and plot the results for inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = fft(x)\n",
    "Xnorm = np.abs(X)/N\n",
    "\n",
    "stem_plot(x      = [n[0:N//2+1]*fs/N/1e3],\n",
    "          y      = [Xnorm[0:N//2+1]],\n",
    "          xticks = [n[0:N//2+1]*fs/N/1e3],\n",
    "          figsize= (6, 3),\n",
    "          title  = [\"Normalise Magnitude Spectra of $X(k)$\"], \n",
    "          xlabel = [\"Frequency (kHz)\"], \n",
    "          ylabel = [\"Amplitude\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, we can compare the FFT with the DFT and find that the results are the exact same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all(np.round(fft(x), 2) == np.round(dft(x), 2))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4. FFT versus DFT <a class=\"anchor\" id=\"fft-dft\"></a>\n",
    "The time required to compute the FFT is significantly faster than the DFT. In this section, we will compare each of our implementations in terms of processing time. Firstly, let us define a test waveform that is 2048 samples long, which means the DFT length is 2048."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 2048    # Number of samples\n",
    "n = np.arange(N) # Sample points\n",
    "\n",
    "a = 0.5*np.cos(2*np.pi*6e3*n*ts + np.pi/4)\n",
    "b = np.cos(2*np.pi*12e3*n*ts)\n",
    "\n",
    "x = np.around(a + b, 4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, compute the execution time of the DFT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "Xdft = dft(x)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For comparison, compute the execution time of the FFT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "Xfft = fft(x)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wow! That was fast! Let us compare both results to check that they are the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all(np.round(Xfft, 2) == np.round(Xdft, 2))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.5. The Numpy FFT Module <a class=\"anchor\" id=\"numpy-fft-module\"></a>\n",
    "You should also be aware the Numpy Python library contains FFT support. The module that you will use for the FFT is known as `np.fft`. This module contains several methods for implementing FFT operations. The FFT is implemented using `np.fft.fft`. See the example given in the code cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.fft.fft(x)\n",
    "Xnorm = np.abs(X)/N"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At any point, you can run the following command to get help using the Numpy FFT method.\n",
    "```python\n",
    "np.fft.fft?\n",
    "```\n",
    "We can plot the results for analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stem_plot(x      = [n[0:N//2+1:128]*fs/N/1e3],\n",
    "          y      = [Xnorm[0:N//2+1:128]],\n",
    "          xticks = [n[0:N//2+1:128]*fs/N/1e3],\n",
    "          figsize= (6, 3),\n",
    "          title  = [\"Normalise Magnitude Spectra of $X(k)$\"], \n",
    "          xlabel = [\"Frequency (kHz)\"], \n",
    "          ylabel = [\"Amplitude\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may also be interested in exploring other methods in the `np.fft` module. These are provided below.\n",
    "\n",
    "```python\n",
    "np.fft.ifft? # Inverse FFT\n",
    "np.fft.fftshift? # FFT shift for swapping each half of the FFT\n",
    "np.fft.fftfreq? # Get the corresponding frequencies for an FFT frame\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Short-Time Fourier Transform <a class=\"anchor\" id=\"stft\"></a>\n",
    "\n",
    "The FFT gives us a lot of insight into sampled signals, and allows us to make many meaningful measurements. However, what if we were also interested in observing how the frequency changes over time? A type of visualization you will often encounter in spectrum analysers, especially in the context of Software Defined Radio (SDR), is the waterfall plot (or spectrogram).\n",
    "\n",
    "Here we will implement one possible implementation of the Short-time Fourier Transform (STFT), which is essentially chopping up the signal of interest into $L$ number of frames, which are then windowed and passed to an FFT. Later, we can stitch these together into a 2-D image that shows our signal's behaviour over time.\n",
    "\n",
    "First off, we will create test stimulus. Let's sum several sine waves together across a long observation period."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set basic params\n",
    "fs = 1024 # sample rate\n",
    "f1 = 60 # frequency of first signal\n",
    "f2 = 120 # frequency of 2nd signal\n",
    "f3 = 250 # frequency of 3rd signal\n",
    "N  = 512 # fft size\n",
    "\n",
    "L = N*256 # we want 256 frames of 512 FFT windows\n",
    "\n",
    "# create a sampled sinusoid and corresponding frequencies for plotting\n",
    "sine_1 = np.sin(2*np.pi*f1*np.arange(L)/fs)\n",
    "sine_2 = 0.5*np.sin(2*np.pi*f2*np.arange(L)/fs)\n",
    "sine_3 = 0.7*np.sin(2*np.pi*f3*np.arange(L)/fs)\n",
    "\n",
    "sum_of_tones = sine_1 + sine_2 + sine_3"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will reshape our sum of tones signal, so that it contains 256 frames of length 512."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum_of_tones = sum_of_tones.reshape(-1,N)\n",
    "sum_of_tones.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Numpy FFT is applied to all the frames in our array. Once again, we will only be looking at the positive side of the spectrum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute FFT and calculate the PSD\n",
    "X = np.fft.fft(sum_of_tones*np.hamming(N), N)[:,:int(N/2)]\n",
    "Xlog = 10*np.log10(2*np.abs(X)**2/N)\n",
    "\n",
    "# Get corresponding frequencies, take only positive spectrum\n",
    "freqs = np.fft.fftfreq(N,1/fs)[:int(N/2)]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can plot the log-scale power spectrum over time. Since we are plotting stationary signals (sine waves) we expect them to be fairly straight and unchanging over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10,5))\n",
    "ax = plt.axes()\n",
    "im = plt.pcolormesh(freqs, np.arange(int(L/N))/fs, Xlog, vmin=-50, shading='gouraud')\n",
    "ax.set_xlabel('Frequency, Hz')\n",
    "ax.set_ylabel('Time, s')\n",
    "ax.set_title('Waterfall Plot')\n",
    "fig.colorbar(im)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1. STFT of a Chirp <a class=\"anchor\" id=\"stft-of-a-chirp\"></a>\n",
    "\n",
    "We can also perform a frequency sweep by taking the STFT of a chirp signal. A chirp is a type of signal that increases or decreases in frequency over time. Here, we construct an up-chirp, by passing a vector of increasing frequency values to our sine function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set basic params\n",
    "fs = 1024 # sample rate\n",
    "N = 512 # fft size\n",
    "L = N*256 # we want 256 frames of 512 FFT windows\n",
    "\n",
    "# Frequency increases over time\n",
    "f = 0.0019*np.arange(L)\n",
    "\n",
    "# Create a sampled sinusoid and corresponding frequencies for plotting\n",
    "chirp = np.sin(2*np.pi*f*np.arange(L)/fs)\n",
    "\n",
    "# Plot the chirp signal\n",
    "stem_plot(x      = [np.arange(0, 4096, 10)],\n",
    "          y      = [chirp[:4096:10]],\n",
    "          xticks = [np.arange(0, 4096, 1000)],\n",
    "          figsize= (14, 3),\n",
    "          title  = [\"Chirp Signal\"], \n",
    "          xlabel = [\"Samples (n)\"], \n",
    "          ylabel = [\"Amplitude\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Performing the same steps as we did for our sum of sines, we visualize the frequency sweep in the waterfall plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Slice it up into 512-sized chunks\n",
    "chirp = chirp.reshape(-1,N)\n",
    "chirp.shape\n",
    "\n",
    "# Compute FFT and calculate the PSD\n",
    "X = np.fft.fft(chirp*np.hamming(N), N)[:,:int(N/2)]\n",
    "Xlog = 10*np.log10(2*np.abs(X)**2/N/fs)\n",
    "\n",
    "# Get corresponding frequencies, take only positive spectrum\n",
    "freqs = np.fft.fftfreq(N,1/fs)[:int(N/2)]\n",
    "\n",
    "fig = plt.figure(figsize=(10,5))\n",
    "ax = plt.axes()\n",
    "im = plt.pcolormesh(freqs, np.arange(int(L/N))/fs, Xlog, vmin=-50, shading='gouraud')\n",
    "ax.set_xlabel('Frequency, Hz')\n",
    "ax.set_ylabel('Time, s')\n",
    "ax.set_title('Waterfall Plot')\n",
    "fig.colorbar(im)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This simple example demonstrates the functionality of a waterfall plot in observing the spectra of a discrete waveform over time. As shown, the frequency of the waveform increases as a function of time by observing the amplitude of the log-scale magnitude."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Conclusion <a class=\"anchor\" id=\"conclusion\"></a>\n",
    "This notebook has presented the FFT. We first investigates the Danielson Lanczos Lemma, which describes how an $N$-point DFT can be computed using two smaller $N/2$-point DFTs. Then we explored an implementation of the FFT and the Numpy FFT using an example. Lastly, we briefly investigated waterfall (spectrogram) plots and their functionality.\n",
    "\n",
    "In the next notebook, we begin learning about digital communications with baseband modulation."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "[⬅️ Previous Notebook](02_investigating_the_dft.ipynb) || [Next Notebook 🚀](../notebook_D/01_baseband_modulation.ipynb)\n",
    "\n",
    "Copyright © 2023 Strathclyde Academic Media\n",
    "\n",
    "---\n",
    "---"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
